# coding: utf8
# Copyright (c) 2019 PaddlePaddle Authors. All Rights Reserve.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#-------------------------------------------------------------
# 注意，数据集路径中不要有中文。已经用 Paddlex 2.1.0 验证过
#-------------------------------------------------------------
from __future__ import print_function

import argparse
import glob
import math
import json
import os
import os.path as osp
import numpy as np
import PIL.Image
import PIL.ImageDraw
import cv2
import math
#from gray2pseudo_color import get_color_map_list


#-------------------
# readfile
#-------------------
def readfile(filename):
    f = open(filename,'rb')
    fs = f.read()
    fs = fs.decode('utf-8')
    f.close()
    return fs

def writefile(filename,content):
    f = open(filename,'wb')
    if type(content) == str:
        content_bytes = content.encode('utf-8')
        f.write(content_bytes)
    else:
        f.write(content)
    f.close()
    return

#-------------------
# copyfile
#-------------------
def copyfile(filename,filename2):
    f = open(filename,'rb')
    fs = f.read()
    f.close()

    f2 = open(filename2,'wb')
    f2.write(fs)
    f2.close()
    

#-------------------
# writefile2
#-------------------
def writefile2(filename,content):
    f = open(filename,'ab')
    if type(content) == str:
        content_bytes = content.encode('utf-8')
        f.write(content_bytes)
    else:
        f.write(content)
    f.flush()
    f.close()
    return
    
#-------------------
# has_file
#-------------------
def has_file(filename):
    if os.path.exists(filename):
        return True
    else:
        return False
        
#-------------------
# remove_file
#-------------------
def remove_file(filename):
    if os.path.exists(filename):
        os.remove(filename)
 
#-------------------
# get_file_size
#-------------------        
def get_file_size(fname):
    if has_file(fname):
        statinfo = os.stat(fname)
        return statinfo.st_size
    else:
        return 0



def get_color_map_list(num_classes):
    """
    Returns the color map for visualizing the segmentation mask,
    which can support arbitrary number of classes.
    Args:
        num_classes (int): Number of classes.
    Returns:
        (list). The color map.
    """

    num_classes += 1
    color_map = num_classes * [0, 0, 0]
    for i in range(0, num_classes):
        j = 0
        lab = i
        while lab:
            color_map[i * 3] |= (((lab >> 0) & 1) << (7 - j))
            color_map[i * 3 + 1] |= (((lab >> 1) & 1) << (7 - j))
            color_map[i * 3 + 2] |= (((lab >> 2) & 1) << (7 - j))
            j += 1
            lab >>= 3
    color_map = color_map[3:]
    return color_map


def gray2pseudo_color(args):
    """将灰度标注图片转换为伪彩色图片"""
    input = args.dir_or_file
    output_dir = args.output_dir
    if not osp.exists(output_dir):
        os.makedirs(output_dir)
        print('Creating colorful label directory:', output_dir)

    color_map = get_color_map_list(256)
    if os.path.isdir(input):
        for fpath, dirs, fs in os.walk(input):
            for f in fs:
                try:
                    grt_path = osp.join(fpath, f)
                    _output_dir = fpath.replace(input, '')
                    _output_dir = _output_dir.lstrip(os.path.sep)

                    im = Image.open(grt_path)
                    lbl = np.asarray(im)

                    lbl_pil = Image.fromarray(lbl.astype(np.uint8), mode='P')
                    lbl_pil.putpalette(color_map)

                    real_dir = osp.join(output_dir, _output_dir)
                    if not osp.exists(real_dir):
                        os.makedirs(real_dir)
                    new_grt_path = osp.join(real_dir, f)

                    lbl_pil.save(new_grt_path)
                    print('New label path:', new_grt_path)
                except:
                    continue
    elif os.path.isfile(input):
        if args.dataset_dir is None or args.file_separator is None:
            print('No dataset_dir or file_separator input!')
            sys.exit()
        with open(input) as f:
            for line in f:
                parts = line.strip().split(args.file_separator)
                grt_name = parts[1]
                grt_path = os.path.join(args.dataset_dir, grt_name)

                im = Image.open(grt_path)
                lbl = np.asarray(im)

                lbl_pil = Image.fromarray(lbl.astype(np.uint8), mode='P')
                lbl_pil.putpalette(color_map)

                grt_dir, _ = osp.split(grt_name)
                new_dir = osp.join(output_dir, grt_dir)
                if not osp.exists(new_dir):
                    os.makedirs(new_dir)
                new_grt_path = osp.join(output_dir, grt_name)

                lbl_pil.save(new_grt_path)
                print('New label path:', new_grt_path)
    else:
        print('It\'s neither a dir nor a file')






def parse_args():
    parser = argparse.ArgumentParser(
        formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument('input_dir', help='input annotated directory')
    return parser.parse_args()




#---------------------
# step1
# 转出annotation图片
#---------------------
def step1():
    

    output_dir = ('Annotations')
    if not osp.exists(output_dir):
        os.makedirs(output_dir)
        print('Creating Annotations directory:', output_dir)

    # get the all class names for the given dataset
    class_names = ['_background_']
    for label_file in glob.glob(osp.join(args.input_dir, '*.json')):
        with open(label_file) as f:
            data = json.load(f)
            for shape in data['shapes']:
                label = shape['label']
                cls_name = label
                if not cls_name in class_names:
                    class_names.append(cls_name)

    class_name_to_id = {}
    for i, class_name in enumerate(class_names):
        class_id = i  # starts with 0
        class_name_to_id[class_name] = class_id
        if class_id == 0:
            assert class_name == '_background_'
    class_names = tuple(class_names)
    print('class_names:', class_names)

    out_class_names_file = osp.join(args.input_dir, 'class_names.txt')
    with open(out_class_names_file, 'w') as f:
        f.writelines('\n'.join(class_names))
    print('Saved class_names:', out_class_names_file)

    color_map = get_color_map_list(256)


    image_file_names = []
    for label_file in glob.glob(osp.join(args.input_dir, '*.json')):
        print('Generating dataset from:', label_file)
        with open(label_file) as f:
            base = osp.splitext(osp.basename(label_file))[0]
            out_png_file = osp.join(output_dir, base + '.png')

            data = json.load(f)

            img_file = osp.join(osp.dirname(label_file), data['imagePath'])
            img = np.asarray(cv2.imread(img_file))

            print(img_file)
            #print('img_file{}, shapes:{}:',img_file,data['shapes'])
            lbl = shape2label(
                img_size=img.shape,
                shapes=data['shapes'],
                class_name_mapping=class_name_to_id,
            )

            if osp.splitext(out_png_file)[1] != '.png':
                out_png_file += '.png'
            # Assume label ranges [0, 255] for uint8,
            if lbl.min() >= 0 and lbl.max() <= 255:
                lbl_pil = PIL.Image.fromarray(lbl.astype(np.uint8), mode='P')
                lbl_pil.putpalette(color_map)
                lbl_pil.save(out_png_file)
                image_file_names.append(img_file)
            else:
                raise ValueError(
                    '[%s] Cannot save the pixel-wise class label as PNG. '
                    'Please consider using the .npy format.' % out_png_file)
                    
    return image_file_names,class_names


#--------------------
# main
#--------------------
def main(args):
    image_file_names, class_names = step1()

    #print('------------------------------------')   
    #print('step 1 executed ....................')
    #print('------------------------------------')
    #print('image_file_names:',image_file_names)
    #print('class_names:',class_names)
    
    n = len(image_file_names)
    # train
    n1 = int(n * 0.7)
    
    # val
    n2 = int(n * 0.2)
    
    # test
    n3 = n - n1 - n2
    
    # label.txt
    remove_file('labels.txt')
    remove_file('train_list.txt')
    remove_file('test_list.txt')
    remove_file('val_list.txt')
    remove_file('train.txt')
    remove_file('test.txt')
    remove_file('val.txt')
    
    if has_file('JPEGImages') == False:
        os.mkdir('JPEGImages')

    s = ''
    for name in class_names:
        s += name + '\r\n'
    s = s.rstrip('\r\n')
    writefile('labels.txt',s)    
    
    print('labels.txt generated')

    s = ''
    for fname in image_file_names[:n1]:
        new_file_path = fname.replace(args.input_dir,'JPEGImages')
        copyfile(fname,new_file_path)
        print('check filename:',new_file_path,cv2.imread(new_file_path).shape)
        s += fname.replace(args.input_dir,'JPEGImages')  + ' ' + fname.replace(args.input_dir,'Annotations').replace('.jpg','.png') + '\r\n'
    s = s.rstrip('\r\n').replace('\\','/')
    writefile('train_list.txt',s)
    writefile('train.txt',s)
    
    print('train_list.txt generated')

    s = ''
    for fname in image_file_names[n1:n1+n2]:
        new_file_path = fname.replace(args.input_dir,'JPEGImages')
        copyfile(fname,new_file_path)
        print('check filename:',new_file_path,cv2.imread(new_file_path).shape)        
        
        s += fname.replace(args.input_dir,'JPEGImages')  + ' ' + fname.replace(args.input_dir,'Annotations').replace('.jpg','.png') + '\r\n'
    s = s.rstrip('\r\n').replace('\\','/')
    writefile('val_list.txt',s)
    writefile('val.txt',s)
    
    print('val_list.txt generated')
    
    s = ''
    for fname in image_file_names[:-n3]:
        new_file_path = fname.replace(args.input_dir,'JPEGImages')
        copyfile(fname,new_file_path)
        print('check filename:',new_file_path,cv2.imread(new_file_path).shape)
        s += fname.replace(args.input_dir,'JPEGImages')  + ' ' + fname.replace(args.input_dir,'Annotations').replace('.jpg','.png') + '\r\n'
    s = s.rstrip('\r\n').replace('\\','/')
    writefile('test_list.txt',s)
    writefile('test.txt',s)
    
    print('test.txt generated')


def shape2mask(img_size, points):
    label_mask = PIL.Image.fromarray(np.zeros(img_size[:2], dtype=np.uint8))
    image_draw = PIL.ImageDraw.Draw(label_mask)
    points_list = [tuple(point) for point in points]
    
    #print('points_list:',points_list)
    
    assert len(points_list) > 2, 'Polygon must have points more than 2'
    image_draw.polygon(xy=points_list, outline=1, fill=1)
    return np.array(label_mask, dtype=bool)

def circle2polygon(points):
    r = math.sqrt((points[0][0]-points[0][1])*(points[0][0]-points[0][1]) + (points[0][1] - points[1][1])*(points[0][1] - points[1][1]))
    x = points[0][0]
    y = points[0][1]
    points2 = []
    for i in range(360):
        x2 = x + r * math.cos(2*3.1415925*i/360.0)
        y2 = y + r * math.sin(2*3.1415925*i/360.0)
        points2.append((x2,y2))
        
    #print('x,y,r:',x,y,r)
    #print('points2:',points2)
    return points2

def shape2label(img_size, shapes, class_name_mapping):
    label = np.zeros(img_size[:2], dtype=np.int32)
    
    print('shape2label shapes:',shapes)
    for shape in shapes:
        points = shape['points']
        class_name = shape['label']
        shape_type = shape.get('shape_type', None)
        
        #print('shape_type:',shape_type)
        
        if shape_type == 'polygon':
            class_id = class_name_mapping[class_name]
            label_mask = shape2mask(img_size[:2], points)
            label[label_mask] = class_id
        elif shape_type == 'circle':
        
            class_id = class_name_mapping[class_name]
            
            points2 = circle2polygon(points)
                        
            
            label_mask = shape2mask(img_size[:2], points2)
            label[label_mask] = class_id
            
    return label


if __name__ == '__main__':
    args = parse_args()
    main(args)
